package top.cardone.data.jdbc.datasource.looku.impl;

import com.google.common.collect.Lists;
import lombok.Synchronized;
import org.springframework.jdbc.datasource.lookup.AbstractRoutingDataSource;
import org.springframework.util.CollectionUtils;
import top.cardone.data.jdbc.datasource.looku.RoutingDataSource;

import java.util.List;
import java.util.Map;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * @author yao hai tao
 */
public class RoutingDataSourceImpl extends AbstractRoutingDataSource implements RoutingDataSource {
	private AtomicInteger counter = new AtomicInteger(1);
	/**
	 * keyListThreadLocal
	 */
	private ThreadLocal<List<Object>> keyListThreadLocal = new ThreadLocal<List<Object>>() {
		@Override
		protected List<Object> initialValue() {
			return Lists.newArrayList();
		}
	};

	private Object[] targetDataSourcesKeys;

	@Override
	public void addKey(Object key) {
		keyListThreadLocal.get().add(key);
	}

	@Override
	protected Object determineCurrentLookupKey() {
		if (CollectionUtils.isEmpty(keyListThreadLocal.get())) {
			return null;
		}

		//按照顺序选择从库
		int index = counter.incrementAndGet() % targetDataSourcesKeys.length;

		if (index < 0) {
			index = -index;
		}

		return targetDataSourcesKeys[index];
	}

	@Synchronized
	@Override
	public void removeKey() {
		int removeIndex = keyListThreadLocal.get().size() - 1;

		if (removeIndex < 1) {
			keyListThreadLocal.remove();

			return;
		}

		keyListThreadLocal.get().remove(removeIndex);
	}

	@Override
	public void removeKeyList() {
		keyListThreadLocal.remove();
	}

	@Override
	public void setTargetDataSources(Map<Object, Object> targetDataSources) {
		super.setTargetDataSources(targetDataSources);

		targetDataSourcesKeys = targetDataSources.keySet().toArray();
	}
}